#!/usr/bin/python3
# python3 extract_origin_data.py

import os, sys, json

def ext_target_dir(source_data_path, dir_t1, dir_rest, user_list, extract_dir = "GU", rest_sub_dir = "FunImg"):
    print('#' * 50)
    print('Start Checking File...')

    mk_dir_list = {}
    # 输出所有文件和文件夹
    for sub_dir in os.listdir(source_data_path):
        # 处理T1数据
        if sub_dir == dir_t1:
            input_dir = os.path.join(source_data_path, sub_dir)

            mk_dir_list[sub_dir] = []
            for user_sub_dir in os.listdir(input_dir):
                if user_sub_dir in user_list:
                    user_file_path = os.path.join(input_dir, user_sub_dir)
                    print('[EXTRACT DIR_NAME]: %s' % (user_file_path))
                    #
                    mk_dir_list[sub_dir].append(user_file_path)

        # 处理rest数据
        elif sub_dir == dir_rest:
            input_dir = os.path.join(source_data_path, sub_dir)
            input_dir = os.path.join(input_dir, extract_dir)
            input_dir = os.path.join(input_dir, rest_sub_dir)

            mk_dir_list[sub_dir] = []
            for user_sub_dir in os.listdir(input_dir):
                if user_sub_dir in user_list:
                    user_file_path = os.path.join(input_dir, user_sub_dir)
                    print('[EXTRACT DIR_NAME]: %s' % (user_file_path))
                    #
                    mk_dir_list[sub_dir].append(user_file_path)

    print("# Total T1 Files: %d" % (len(mk_dir_list[dir_t1])))
    print("# Total rest Files: %d" % (len(mk_dir_list[dir_rest])))
    print('Checking File Finished!')
    print('#' * 50)
    return mk_dir_list

def read_user_list(file_path):
    user_list = []

    fp = open(file_path, "r")
    for line in fp.readlines():
        line_list = line.split()
        item_name = line_list[0] + line_list[1] + '_' + line_list[3]

        user_list.append(item_name)

    return user_list

def cp_file_to_target(dirs_dict, target_dirs, dir_t1, dir_rest):
    target_t1_dir = os.path.join(target_dirs, dir_t1)
    target_rest_dir = os.path.join(target_dirs, dir_rest)

    if not os.path.exists(target_t1_dir):
        os.makedirs(target_t1_dir)

    if not os.path.exists(target_rest_dir):
        os.makedirs(target_rest_dir)

    print('#' * 50)
    print('Start Copying File...')

    for (dir_name, dir_list) in dirs_dict.items():
        if dir_name == dir_t1:
            cmd_str = ""
            for dir_t1_item in dir_list:
                cmd_str = "cp -dR " + dir_t1_item + " " + target_t1_dir

                print("[CMD]: %s" % (cmd_str))
                # os.system(cmd_str)
        if dir_name == dir_rest:
            cmd_str = ""
            for dir_rest_item in dir_list:
                sub_symbol = dir_rest_item.split('/')[-1]
                sub_id = sub_symbol.split('_')[-1]

                if "OHSU" in sub_symbol:
                    for OHSU_dir in ['/home/li/Documents/MRI/test-retest/OHSU_1/s36', '/home/li/Documents/MRI/test-retest/OHSU_2/s36']:
                        for ori_rest_dir in ['rest_1', 'rest_2', 'rest_3']:
                            sess_target_rest_dir = os.path.join(target_rest_dir, sub_symbol + '/' + ori_rest_dir)
                            if not os.path.exists(sess_target_rest_dir):
                                os.makedirs(sess_target_rest_dir)

                            ori_rest_sess_dir = os.path.join(OHSU_dir, ori_rest_dir)
                            ori_rest_list = os.listdir(ori_rest_sess_dir)

                            if sub_id in ori_rest_list:
                                src_rest_sess_dir = os.path.join(ori_rest_sess_dir, sub_id)

                                cmd_str = "cp -dR " + src_rest_sess_dir + "/* " + sess_target_rest_dir + '/'
                                print("[CMD]: %s" % (cmd_str))
                                # os.system(cmd_str)
                else:
                    sess_target_rest_dir = os.path.join(target_rest_dir, sub_symbol + '/rest_1')
                    if not os.path.exists(sess_target_rest_dir):
                        os.makedirs(sess_target_rest_dir)

                    cmd_str = "cp -R " + dir_rest_item + "/* " + sess_target_rest_dir

                    print("[CMD]: %s" % (cmd_str))
                    # os.system(cmd_str)

    print('Copying File Finished!')
    print('#' * 50)

def write_recon_list(dirs_dict, recon_file, dir_t1):
    print('#' * 50)
    print("Writing `recon_list.txt` to `%s`" % (recon_file))

    line_list = [line_item.split('/')[-1] + "\n" for line_item in dirs_dict[dir_t1]]

    fp = open(recon_file, "w")
    fp.writelines(line_list)
    # 关闭文件
    fp.close()

    print("Done!")
    print('#' * 50)

def tmp_write_recon_list(user_list, recon_file):
    print('#' * 50)
    print("Writing `recon_list.txt` to `%s`" % (recon_file))

    line_list = [line_item + "\n" for line_item in user_list]

    fp = open(recon_file, "w")
    fp.writelines(line_list)
    # 关闭文件
    fp.close()

    print("Done!")
    print('#' * 50)

if __name__ == '__main__':
    symbol = "OHSU2"
    extract_dir = "OHSU"

    # 用户列表文件
    user_file = './user_list_' + symbol + '.txt'
    # 源文件地址
    source_data_path = '/home/li/Documents/MRI/ASD_data'
    # 目标文件地址
    target_data_path = '/mri_projects/ASD_Analysis/Origin'
    # recon 参考文件
    recon_file = '/mri_projects/ASD_Analysis/Lists/recon_list_' + symbol + '.txt'

    dir_t1 = 'T1'
    dir_rest = 'rest'

    user_list = read_user_list(user_file)
    dir_dict = ext_target_dir(source_data_path, dir_t1, dir_rest, user_list, extract_dir)
    cp_file_to_target(dir_dict, target_data_path, dir_t1, dir_rest)
    write_recon_list(dir_dict, recon_file, dir_t1)
    # tmp_write_recon_list(user_list, recon_file)
